import os
import json
import time
import pickle
import numpy as np
import pandas as pd
from multiprocessing import Pool, cpu_count
from itertools import repeat
from ensemble_boxes import *
from map_boxes import *
from tqdm import  tqdm

class NpEncoder(json.JSONEncoder):

    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)

def load_annotations(anno_dir):
    '''
    :param anno_dir:
    :return: data -> 'annotations', 'categories', 'images'
    images: {file_name': 'pure_bg_images/bb_V0033_I0004040.jpg', 'height': 1080, 'width': 1920, 'id': 793}
    '''
    with open(anno_dir) as fp:
        data = json.load(fp)
    cat_df = pd.DataFrame(data['categories'])
    anno_df = pd.DataFrame(data['annotations'])
    img_df = pd.DataFrame(data['images'])

    return anno_df, img_df, cat_df

# def get_single_id(anno_df, img_id, height, width):
#     bboxes = []
#
#     # boxes_list
#     df = anno_df[anno_df.image_id == img_id]
#     for i in range(df.shape[0]):
#         x1, y1, w, h = df.bbox.iloc[i]
#         x2 = x1 + w
#         y2 = y1 + h
#
#         bboxes.append([x1/width, y1/height, x2/width, y2/height])
#         # bboxes.append([x1, y1, x2, y2])
#     # scores_list
#     scores = df.score.to_list()
#     # labels_list
#     labels = df.category_id.to_list()
#
#     return list(bboxes), scores, labels

def get_single_id(anno_df, img_id, height, width):
    bboxes = []
    scores = []
    labels = []

    # boxes_list
    df = anno_df[anno_df.image_id == img_id]
    for i in range(df.shape[0]):
        x1, y1, w, h = df.bbox.iloc[i]
        x2 = x1 + w
        y2 = y1 + h
        score = df.score.iloc[i]
        cat = df.category_id.iloc[i]
        if w * h > 784 or w*h<2:
            continue
        bboxes.append([x1 / width, y1 / height, x2 / width, y2 / height])
        scores.append(score)
        labels.append(cat)

    return list(bboxes), list(scores), list(labels)

def ensemble_single_img(anno_dfs, height, width, img_id, weights=[], iou_thr=0.5, iou_type='iou', conf_type='avg'):
    boxes_list = []
    scores_list = []
    labels_list = []

    for anno_df in anno_dfs:
        bbox, score, label = get_single_id(anno_df, img_id, height, width)
        boxes_list.append(bbox)
        scores_list.append(score)
        labels_list.append(label)

    norm_bboxes, scores, labels = weighted_boxes_fusion(boxes_list,
                                                        scores_list,
                                                        labels_list,
                                                        weights=weights,
                                                        allows_overflow=False,
                                                        iou_thr=iou_thr,
                                                        skip_box_thr=0.0,
                                                        conf_type=conf_type,
                                                        iou_type=iou_type)

    bboxes = []
    for bbox in norm_bboxes:
        x1, y1, x2, y2 = bbox
        x1 = x1 * width
        y1 = y1 * height
        x2 = x2 * width
        y2 = y2 * height
        w = x2 - x1
        h = y2 - y1
        bboxes.append([x1, y1, w, h])

    return bboxes, scores, labels

def ensemble(anno_dfs, img_org_df, weights=[], iou_thr=0.5, iou_type='iou', conf_type='avg'):
    annotations = []
    for i in tqdm(range(img_org_df.shape[0])):
        img_id = img_org_df.id.iloc[i]
        height = img_org_df.height.iloc[i]
        width = img_org_df.width.iloc[i]
        bboxes, scores, labels = ensemble_single_img(anno_dfs,
                                                     height,
                                                     width,
                                                     img_id,
                                                     weights=weights,
                                                     iou_thr=iou_thr,
                                                     iou_type=iou_type,
                                                     conf_type=conf_type)
        for j in range(len(bboxes)):
            x, y, w, h = bboxes[j]
            # if w*h > 1000 and scores[j] < 0.05:
            #     continue
            annotations.append({'image_id': img_id,
                                'bbox': bboxes[j],
                                'score': scores[j],
                                'category_id': 1})
    res_dir = 'last_new2.json'
    with open(res_dir, 'w') as fp:
        # json.dump(annotations, fp, indent=4, separators=(',', ': '), cls=NpEncoder)
        json.dump(annotations, fp, cls=NpEncoder)


def main():
    # anno_listdir = [
    #     'C:\\Work\\competition\\tiny\\src\\fr50_re2net_visdrone_mst.json', # 64.08
    #     'C:\\Work\\competition\\tiny\\src\\fr50_06166_mst.json', # 63.53
    #     'C:\\Work\\competition\\tiny\\src\\fr50_re2net_visdrone_cutmix_200.json', #64.65
    #     'C:\\Work\\competition\\tiny\\src\\fr50_resnet_visdrone_mst.json', #63.36
    #     'C:\\Work\\competition\\tiny\\src\\cb50_visdrone_200.json', # 62.610
    #     'C:\\Work\\competition\\tiny\\src\\fr50_resnet_visdrone_cutmix_mst.json', # 64.45
    #     'C:\\Work\\competition\\tiny\\src\\fr101.json',
    #     'C:\\Work\\competition\\tiny\\src\\frxt101.json',
    #     'C:\\Work\\competition\\tiny\\src\\cr50_dcn.json',
    #     # 'C:\\Work\\competition\\tiny\\src\\sub_tint_fr_regnet_dcn_smoothl22.json',
    #     'C:\\Work\\competition\\tiny\\src\\sub_tiny_fr_resnet_dcn_fpn_bn_large.json',
    #     'C:\\Work\\competition\\tiny\\src\\sub_tiny_fr_regent_dcn_fpn_bn_large.json',
    #     'C:\\Work\\competition\\tiny\\src\\sub_tiny_fr_hrnet2p_w18_base.json'
    # ]


# new wbf
# 68.37
# anno_listdir = [
#     'fr_res2net50_anchor3_1024.json',
#     'fr50_resnet_visdrone_cutmix_fpn_bn.json',
#     ]
    anno_listdir = [
        'fr_res2net50_anchor3_1024.json',
        'fr50_resnet_visdrone_cutmix_fpn_bn.json', # 68.37
        'cb50_fpn_bn_anchor3.json', # 68.63
        'sub_tiny_fr_hrnet2p_w18_base.json', # 69.52
        'cascade_r50_dcn_anchor3_fpn_bn.json',
        'fr101.json',
        'fr101_fpn_bn_anchor3.json',
        'frxt101_dcn_anchor3_fpn_bn.json',
        'frxt101_1024.json'
    ]

    anno_dfs = []
    for anno_dir in anno_listdir:
        anno_dfs.append(pd.read_json('C:\\Work\\competition\\tiny\\src\\wbf\\'+anno_dir))
    # weights = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,1]
    weights = [1] * len(anno_listdir)
    iou_thr = 0.58

    iou_type = 'iou'
    conf_type = 'avg'
    test_org_anno_dir = 'C:\\Work\\competition\\tiny\\src\\tiny_set_test_nobox.json'
    _, img_org_df, _ = load_annotations(test_org_anno_dir)
    ensemble(anno_dfs, img_org_df, weights=weights, iou_thr=iou_thr, iou_type=iou_type, conf_type=conf_type)

if __name__ == '__main__':
    main()
